#include "StdAfx.h"
#include "Loader.h"
#include "SetInfo.h"
#include "Visitor.h"

namespace MnistDigit
{
	Loader::Loader(const SetInfo& setInfo, Visitor& visitor)
		: m_labelSet(setInfo.m_labelFileName.c_str(), std::ios::in | std::ios::binary)
		, m_imageSet(setInfo.m_imageFileName.c_str(), std::ios::in | std::ios::binary)
		, m_visitor(visitor)
		, m_setInfo(setInfo)
	{
		if(!m_labelSet.is_open() || !m_labelSet.good())
		{
			std::string message;
			message.assign(setInfo.m_labelFileName.begin(), setInfo.m_labelFileName.end()); 
			throw std::logic_error(message+ "- Can not open");
		}

		if(!m_imageSet.is_open() || !m_labelSet.good())
		{
			std::string message;
			message.assign(setInfo.m_imageFileName.begin(), setInfo.m_imageFileName.end()); 
			throw std::logic_error(message+ "- Can not open");
		}
	}

	Loader::~Loader(void)
	{
	}

	void Loader::LoadSets()
	{
		LoadHeaders();
		LoadData();
	}

	void Loader::LoadData()
	{
		size_t dataCount = 0;
		unsigned char digit=0;
		const size_t imageSize = m_imageHeader.number_of_rows*m_imageHeader.number_of_columns;
		std::vector<unsigned char> dataBuffer(imageSize);

		while(m_labelSet.read( reinterpret_cast<char*>(&digit), sizeof(digit) )
			&& m_imageSet.read( reinterpret_cast<char*>(&dataBuffer[0]), imageSize))
		{
			AddDigit(digit, dataBuffer);
			++dataCount;

			if(m_setInfo.m_dataCountLimit 
				&& dataCount >= m_setInfo.m_dataCountLimit)
				break;
		}

		if(dataCount!=m_imageHeader.m_number_of_images 
			&& dataCount!=m_setInfo.m_dataCountLimit)
			throw std::logic_error("Loader::LoadData");
	}

	void Loader::AddDigit(const unsigned char digit, const std::vector<unsigned char>& dataBuffer)
	{
		m_visitor.Visit(digit, dataBuffer);
	}

	void Loader::LoadHeaders()
	{
		LoadLabelHeader();
		LoadImageHeader();
		if(m_imageHeader.m_number_of_images!=m_labelHeader.m_number_of_items)
			throw std::logic_error("Loader::LoadHeaders");
	}

	void Loader::LoadLabelHeader()
	{
		m_labelSet>>m_labelHeader;
	}

	void Loader::LoadImageHeader()
	{
		m_imageSet>>m_imageHeader;
	}
}